package xjggb.algorithm.树.红黑树;

import java.util.Comparator;

public class AVLTree<E> extends BST<E> {

    public AVLTree(){
        this(null);
    }

    public AVLTree(Comparator<E> comparator){
        super(comparator);
    }


    @Override
    protected void afterAdd(Node<E> node) {
        while ((node=node.parent)!=null){
            if (isBalanced(node)){
                //更新高度
                updateHeight(node);
            }else {
                rebalance( node);
                break;
            }
        }

    }




    /*
    * 恢复平衡
    * */
    private void rebalance(Node<E> grand){
        Node<E> parent = ((AVLNode<E>)grand).tallerChild();
        Node<E> node = ((AVLNode<E>)parent).tallerChild();

        if (parent.isLeftChild()){ //L
            if (node.isLeftChild()){ //LL
                rotateRight(grand);
            }else { //LR
                 rotateLeft(parent);
                 rotateRight(grand);
            }
        }else { //R
            if (node.isLeftChild()){//RL
                rotateRight(parent);
                rotateLeft(grand);
            }else { //RR
                rotateLeft(grand);

            }
        }
    }
    /*
    * 左旋转
    * */
    private void rotateLeft(Node<E> grand){
        Node<E> parent = grand.right;
        Node<E> child = parent.left;
        grand.right = child;
        parent.left = grand;
        afterRotate(grand, parent, child);
    }
    /*
    * 又旋转
    * */
    private void rotateRight(Node<E> grand){
        Node<E> parent = grand.left;
        Node<E> child = parent.right;
        grand.left = child;
        parent.right = grand;
        afterRotate(grand, parent, child);
    }

    private void afterRotate(Node<E> grand,Node<E> parent,Node<E> child){
        // 让parent称为子树的根节点
        parent.parent = grand.parent;
        if (grand.isLeftChild()) {
            grand.parent.left = parent;
        } else if (grand.isRightChild()) {
            grand.parent.right = parent;
        } else { // grand是root节点
            root = parent;
        }

        // 更新child的parent
        if (child != null) {
            child.parent = grand;
        }
        // 更新grand的parent
        grand.parent = parent;

        // 更新高度
        updateHeight(grand);
        updateHeight(parent);
    }




    private boolean isBalanced(Node<E> node) {
        return Math.abs(((AVLNode<E>)node).balanceFactor()) <= 1;
    }



    /*
    * 更新高度
    * */
    private void updateHeight(Node<E> node) {
        ((AVLNode<E>)node).updateHeight();
    }


    /*
    *添加节点高度
    * */
    private static  class AVLNode<E> extends Node<E>{
        int height=1;  //设置平衡因子
        public AVLNode(E element, Node<E> parent) {
            super(element, parent);
        }

        /*
        * 提供平衡因子
        * */
        public  int balanceFactor(){
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            return leftHeight - rightHeight;
        }

        public void updateHeight(){
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            height = 1 + Math.max(leftHeight, rightHeight);
        }

        public Node<E> tallerChild(){
            int leftHeight =left==null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight =right==null ? 0 : ((AVLNode<E>)right).height;
            if (leftHeight > rightHeight) return left;
            if (leftHeight < rightHeight) return right;
            return isLeftChild() ? left : right;
        }
    }

    @Override
    protected void afterRemove(Node<E> node) {
        while ((node = node.parent) != null) {
            if (isBalanced(node)) {
                // 更新高度
                updateHeight(node);
            } else {
                // 恢复平衡
                rebalance(node);
            }
        }
    }

    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new AVLNode<>(element,parent);
    }
}
